import os
import torch
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from torchvision import datasets
from alexnet import AlexNet   # 确保load_model函数已更新
import matplotlib.pyplot as plt
os.environ['KMP_DUPLICATE_LIB_OK']='True'
# 设置设备为GPU/CPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 数据预处理
transform = transforms.Compose([
    transforms.Resize((224, 224)),  # 调整图像大小为224x224
    transforms.ToTensor(),  # 转换为Tensor
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # 归一化
])

# 加载测试数据
test_dir = r'F:\pythonProject\Alexnet\pizza\test'
test_data = datasets.ImageFolder(test_dir, transform=transform)
test_loader = DataLoader(test_data, batch_size=1, shuffle=False)  # 修改batch_size为1以便于可视化

# 初始化模型
num_classes = 3  # 确保这里的类别数与训练时一致
model = AlexNet(num_classes=num_classes).to(device)

# 加载训练好的模型
def load_model(model, path):
    model.load_state_dict(torch.load(path, map_location=device), strict=False)  # 使用strict=False以避免尺寸不匹配错误
    return model

model = load_model(model, r'F:\pythonProject\Alexnet\models\alexnet_best.pth')

# 测试函数
def test():
    model.eval()
    correct = 0
    total = 0

    with torch.no_grad():
        for inputs, labels in test_loader:
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    accuracy = 100 * correct / total
    print(f'Test Accuracy: {accuracy:.2f}%')

    # 可视化
    visualize_samples(test_loader, model, device)

# 可视化测试样本
def visualize_samples(dataloader, model, device):
    model.eval()
    classes = dataloader.dataset.classes
    fig, axs = plt.subplots(3, 3, figsize=(15, 15))
    axs = axs.flatten()
    for i in range(9):
        input, label = next(iter(dataloader))
        input = input.to(device)
        output = model(input)
        _, predicted = torch.max(output, 1)
        axs[i].imshow(input[0].permute(1, 2, 0).cpu().numpy())  # 转换为PIL图像格式
        axs[i].set_title(f"Pred: {classes[predicted.item()]}\nTrue: {classes[label.item()]}", fontsize=12)
        axs[i].axis('off')
    plt.tight_layout()
    plt.savefig('img/test_samples.png')  # 保存可视化图片
    plt.show()

if __name__ == '__main__':
    test()